# encoding=utf-8
import onnxruntime
import torch
from DeeplabModel import DeepLabHead, DeepLabV3, resnet50

# load model
classifier = DeepLabHead(2048, 6)
backbone = resnet50(replace_stride_with_dilation=[False, True, True])
model = DeepLabV3(backbone, classifier, None)

weights_dict = torch.load("model_241.pth", map_location='cpu')['model']
for k in list(weights_dict.keys()):
    if "aux" in k:
        del weights_dict[k]

# load data
err_info = model.load_state_dict(weights_dict, strict=False)
print(err_info)
model.eval()

# in_tensor = torch.randn(1, 3, 512, 512)
# out = model(in_tensor)
# print(out.shape)

# output onnx
torch.onnx.export(
    model,
    torch.randn(1, 3, 512, 512),
    "export_oct_deeplabv3.onnx",
    verbose=True,
    opset_version=16,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output']
)

session = onnxruntime.InferenceSession("export_oct_deeplabv3.onnx")
inputs = {"input":  torch.randn(1, 3, 512, 512).numpy()}
out = session.run(None, inputs)
print(out[0].shape)